/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.server;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;

import org.slf4j.Logger;
import org.sss.common.io.DBSet;
import org.sss.common.io.KeyFormatter;
import org.sss.mapreduce.DB;
import org.sss.mapreduce.Decoder;
import org.sss.mapreduce.Encodings;
import org.sss.mapreduce.GroupID;
import org.sss.mapreduce.SideData;
import org.sss.mapreduce.SssException;
import org.sss.mapreduce.StorageNode;
import org.sss.mapreduce.Tuple;
import org.sss.mapreduce.datatype.Packable;

public class SideDataInAdvance<K extends Packable, V extends Packable> implements SideData<K, V> {
  @SuppressWarnings("unused")
  private static Logger logger = SssServer.getLogger(SideDataInAdvance.class);

  private final Map<K, List<V>> dic = new HashMap<K, List<V>>();

  public static <KEY extends Packable, VALUE extends Packable> SideData<KEY, VALUE> create(
      SssServerResource sssResouce,
      StorageNode node,
      GroupID gid,
      Encodings encodings,
      Class<KEY> keyClass,
      Class<VALUE> valueClass) throws SssException, IOException, ExecutionException, InterruptedException {
    return new SideDataInAdvance<KEY, VALUE>(sssResouce, node, gid, encodings, keyClass, valueClass);
  }

  // TODO: uses appropriate exception.
  // TODO: issue database requests from i/o thread.
  private SideDataInAdvance(SssServerResource sssResouce,
                  StorageNode nearNode,
                  GroupID gid,
                  Encodings encodings,
                  Class<K> keyClass,
                  Class<V> valueClass) throws SssException, IOException, ExecutionException, InterruptedException {
    DBSet dbSet = sssResouce.getDBSetForSideData();
    try {
      SssServerConfiguration sssConf = sssResouce.getServerConfiguration();
      int count_max = sssConf.ioConf.get_batchcount;
      int limitNBytes = sssConf.ioConf.get_limit_nbytes;
      boolean vanilla = sssConf.ioConf.db_use_vanilla;
      boolean destructive = false; // TODO: see config.

      Decoder keyDec = encodings.createDecoder(keyClass);
      Decoder valDec = encodings.createDecoder(valueClass);

      if (gid.isBroadcast()) {
        if (vanilla) {
          restrictedRead(dbSet.get(nearNode), count_max, gid, keyClass,
                         valueClass, keyDec, valDec, dic, destructive);
        }
        else {
          read(dbSet.get(nearNode), limitNBytes, gid, keyClass, valueClass, keyDec, valDec, dic);
        }
      }
      else {
        for (StorageNode node: sssResouce.getStorageNodeList()) {
          if (vanilla) {
            restrictedRead(dbSet.get(node), count_max, gid, keyClass,
                           valueClass, keyDec, valDec, dic, destructive);
          }
          else {
            read(dbSet.get(node), count_max, gid, keyClass, valueClass, keyDec, valDec, dic);
          }
        }
      }
    }
    finally {
      dbSet.close();
    }
  }

  public boolean containsKey(K key) {
    return dic.containsKey(key);
  }

  // TODO: uses appropriate exception.
  public synchronized List<V> get(K key) throws SssException {
    return dic.get(key);
  }

  public Iterable<K> keys() throws SssException {
    return dic.keySet();
  }

  public Iterable<Tuple<K, V>> tuples() throws SssException {
    return new Iterable<Tuple<K, V>>() {
      @Override
      public Iterator<Tuple<K, V>> iterator() {
        return new TupleIterator<K, V>(SideDataInAdvance.this);
      }
    };
  }

  @Override
  public synchronized void close() throws IOException {
    /* Do nothing */
  }

  private static <K extends Packable, V extends Packable> void
    read(DB db,
        int limitNBytes,
        GroupID groupid,
        Class<K> keyClass,
        Class<V> valClass,
        Decoder keyDec, Decoder valDec,
        Map<K, List<V>> container) throws IOException {

    KeyFormatter kf = KeyFormatter.get(groupid.getKeyFormat());
    boolean initial = true;
    byte[] prefix   = KeyFormatter.getGroupIDPrefix(groupid);
    byte[] sentinel = KeyFormatter.getGroupIDSentinel(groupid);
    int count = 0;
    while (true) {
      List<byte[]> kvs = db.range(prefix, sentinel, limitNBytes, initial);
      initial = false;
      count = kvs.size();
      if (count > 0) {
        prefix = kvs.get(count - 2);

        for (int i = 0;i < count;i += 2) {
          K userkey = keyClass.cast(keyDec.decode(kf.getUserKey(kvs.get(i))));
          V value   = valClass.cast(valDec.decode(kvs.get(i+1)));

          List<V> values = container.get(userkey);
          if (values == null) {
            values = new ArrayList<V>();
            container.put(userkey, values);
          }
          values.add(value);
        }
      } else {
        break;
      }
    }
  }

  private static <K extends Packable, V extends Packable> void restrictedRead
    (DB db, int count_max, GroupID groupid, Class<K> keyClass, Class<V> valClass,
     Decoder keyDec, Decoder valDec, Map<K, List<V>> container, boolean destructive)
    throws IOException {
    final KeyFormatter kf = KeyFormatter.get(groupid.getKeyFormat());
    final byte[] prefix = KeyFormatter.getGroupIDPrefix(groupid);
    final List<byte[]> allkeys = db.fwmkeys(prefix, -1);
    final int count_target = allkeys.size();
    boolean done = count_target == 0;
    int count_sofar = 0;
    while (!done) {
      final int count_tmp = Math.min(count_target - count_sofar, count_max);
      byte[][] ks = new byte[count_tmp][];
      for (int i=0; i<count_tmp; i++) {
        ks[i] = allkeys.get(count_sofar + i);
      }
      List<byte[]> kvs = db.getlist(ks);
      final int count_read = kvs.size();
      count_sofar += count_read / 2;
      if (destructive && count_read > 0) {
        db.outlist(ks);
      }
      if (count_read > 0) {
        for (int i=0; i<count_read; i+=2) {
          K userkey = keyClass.cast(keyDec.decode(kf.getUserKey(kvs.get(i))));
          V value   = valClass.cast(valDec.decode(kvs.get(i + 1)));
          List<V> values = container.get(userkey);
          if (values == null) {
            values = new ArrayList<V>();
            container.put(userkey, values);
          }
          values.add(value);
        }
      }
      done = count_target == count_sofar;
    }
  }
}
